#include "Containers/RangePair.h"
#include "File/BOSSOutputLoader.h"
#include "File/YamlLoader.h"
#include "Fit/RooFitBuilder.h"
#include "Functions/Draw.h"
#include "Globals/Settings.h"
#include "Particle/ReconstructedParticle.h"
#include "RooHistPdf.h"
#include "RooPlot.h"
#include "TCanvas.h"
#include "TStyle.h"

using namespace BOSS_Afterburner;
using namespace BOSS_Afterburner::Draw;
using namespace BOSS_Afterburner::Error;
using namespace RooFit;
using namespace BOSS_Afterburner::Detector::Units;
using namespace std;

void PlotDataHistOnFrame(RooDataHist& hist, RooPlot* frame)
{
  hist.plotOn(frame, LineWidth(2), LineColor(kBlue + 2), LineWidth(1), MarkerColor(kBlue + 2),
              MarkerSize(.5));
}

void SaveFrame(RooPlot* frame, const char* name)
{
  TCanvas c;
  frame->Draw();
  SaveCanvas(name, &c);
}

void ExcplicatePeaks(const char* config_file)
{
  YamlLoader conf(config_file);

  // * Data files * //
  const auto&      files = conf["Input files"];
  BOSSOutputLoader excl(files["excl"].as<string>().c_str());
  BOSSOutputLoader incl(files["incl"].as<string>().c_str());
  BOSSOutputLoader data(files["data"].as<string>().c_str());

  // * Cut ranges * //
  const auto& cutvals = conf["cutvals"];
  RangePair   cutD0   = cutvals["D0"];
  RangePair   cutphi  = cutvals["phi"];

  // * Plot and fit ranges * //
  const auto& ranges = conf["Ranges"];
  RangePair   fit    = ranges["fit"];
  RangePair   plot   = ranges["plot"];

  ReconstructedParticle particle(conf["Reconstructed particle"]);

  const char* varname = Form("#it{M}_{%s}", particle.NameLaTeX());
  RooRealVar  var(varname, varname, plot.X1(), plot.X2(), GeVc2);

  RooFitBuilder builder(var);

  BranchPlotOptions exclSigOpts(conf["Exclusive Monte Carlo signal"]);
  BranchPlotOptions inclSigOpts(conf["Inclusive Monte Carlo signal"]);
  BranchPlotOptions dataSigOpts(conf["Measurement signal"]);
  exclSigOpts[0].SetRange(plot.X1(), plot.X2());
  inclSigOpts[0].SetRange(plot.X1(), plot.X2());
  dataSigOpts[0].SetRange(plot.X1(), plot.X2());

  auto exclHist = excl.Draw(exclSigOpts);
  auto inclHist = incl.Draw(inclSigOpts);
  auto dataHist = data.Draw(dataSigOpts);

  RooDataHist rooexclHist("excl", exclHist->GetTitle(), var, Import(*exclHist));
  RooDataHist rooinclHist("incl", inclHist->GetTitle(), var, Import(*inclHist));
  RooDataHist roodataHist("data", dataHist->GetTitle(), var, Import(*dataHist));

  auto exclFrame = var.frame();
  auto inclFrame = var.frame();
  auto dataFrame = var.frame();

  PlotDataHistOnFrame(rooexclHist, exclFrame);
  PlotDataHistOnFrame(rooinclHist, inclFrame);
  PlotDataHistOnFrame(roodataHist, dataFrame);

  SaveFrame(exclFrame, "exclFrame");
  SaveFrame(inclFrame, "inclFrame");
  SaveFrame(dataFrame, "dataFrame");

  // * Construct PDF for hist * //
  RooHistPdf histPdf("histPdf", "histPdf", var, rooexclHist, 0);

  // * Extract parameters from YAML * //
  const auto& variables = conf["Roofit variables"];
  if(!variables) return;
  const auto& varTrunc = variables["Truncated polynomial"];
  const auto& varArgus = variables["Argus"];
  const auto& varArBES = variables["Argus BESIII"];
  if(!varTrunc) return;
  if(!varArgus) return;

  auto ratio = builder.BuildVar(variables["Signal-to-background ratio"]);
  auto argus = builder.BuildArgus(varArgus);
  auto arBES = builder.BuildArgusBESIII(varArBES);
  auto trunc = builder.BuildTruncatedPolynomial(varTrunc);

  // * Add pdfs * //
  auto sum   = new RooArgList;
  auto coeff = new RooArgList;
  sum->add(histPdf);
  auto bckType = variables["Background type"].as<string>();
  bool hasBck  = false;
  if(bckType.compare("Argus") == 0)
  {
    sum->add(*argus);
    hasBck = true;
  }
  else if(bckType.compare("Truncated polynomial") == 0)
  {
    sum->add(*trunc);
    hasBck = true;
  }
  if(hasBck) coeff->add(*ratio);
  auto shape = builder.AddPdfs(sum, coeff);

  // * Perform fit for p-wavve only * //
  Color_t colour = kRed;
  shape->fitTo(roodataHist, Range(fit.X1(), fit.X2()));
  shape->plotOn(dataFrame, LineColor(colour), Range(plot.X1(), plot.X2()), LineWidth(1));
  if(bckType.compare("Argus") == 0)
    shape->plotOn(dataFrame, Components(*argus), LineColor(colour), LineStyle(kDashed),
                  Range(plot.X1(), plot.X2()), LineWidth(1));
  if(bckType.compare("Truncated polynomial") == 0)
    shape->plotOn(dataFrame, Components(*trunc), LineColor(colour), LineStyle(kDashed),
                  Range(plot.X1(), plot.X2()), LineWidth(1));
  shape->paramOn(dataFrame, Layout(.59, .92, .92), Format("NU"));

  // * Configure frame style * //
  if(dataFrame)
  {
    dataFrame->getAttText()->SetTextSize(.035);
    dataFrame->getAttText()->SetTextColor(colour);
  }

  // * Write fitted histograms * //
  TCanvas canvas;
  canvas.SetBatch();
  dataFrame->Draw();
  IndicateFitRange(conf);
  SaveCanvas(Form("ExclToData_%s_%s", particle.Name(), bckType.c_str()), &canvas);
}

/// Main function that is called when executing the executable compiled using e.g. `g++`.
int main(int argc, char* argv[])
{
  TerminalIO::CheckMainArguments(argc, argv, 1);
  ExcplicatePeaks(argv[1]);
  return 0;
}